import logging
from io import BytesIO
from PIL import Image
import numpy as np

import os

os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3'  # 关闭提示打印
logging.warning("正在装载验证码识别模型")

from tensorflow import keras

model = keras.models.load_model("./model/Model_tf.net")


def split_pic(img) -> list:
    img = img.convert('L').convert('1')
    x_size, y_size = img.size  # 72 * 22
    y_size -= 5  # 17
    piece = (x_size - 24) / 8  # 6
    centers = [4 + piece * (2 * i + 1) for i in range(4)]

    ar = []
    for i, center in enumerate(centers):
        single_pic = img.crop(
            (center - (piece + 2), 1, center + (piece + 2), y_size))
        ar.append(np.asarray(single_pic, dtype='int8'))
    return ar


def predict(stream) -> str:
    def func(x): return x + 48 if x <= 9 else x + 87 if x <= 23 else x + 88

    image = Image.open(BytesIO(stream))
    # 先转化为灰度图像,  再转化为[0|1]值图像
    result = []
    data = np.zeros((1, 21, 16), dtype="int8")
    for single in split_pic(image):
        data[0] = single
        answer = model.predict(data)  # 此时answer->[36] 每个值都是介于0~1间 总和为1
        answer = np.argmax(answer)  # 找出预测最有信心的
        result.append(chr(func(answer)))  # 将这个值对应转换成character
    return "".join(result)
